from django.db import migrations

def insert_word_pairs(apps, schema_editor):
    WordPair = apps.get_model('trend', 'WordPair')
    
    pairs_data = [
        # 基础理论关联
        {'kword1_id': 1, 'kword2_id': 2, 'co_count': 800},  # 人工智能-深度学习
        {'kword1_id': 2, 'kword2_id': 4, 'co_count': 600},  # 深度学习-神经网络
        {'kword1_id': 3, 'kword2_id': 2, 'co_count': 500},  # 机器学习-深度学习
        {'kword1_id': 5, 'kword2_id': 2, 'co_count': 400},  # 强化学习-深度学习
        
        # 计算机视觉关联
        {'kword1_id': 6, 'kword2_id': 2, 'co_count': 700},  # 计算机视觉-深度学习
        {'kword1_id': 7, 'kword2_id': 6, 'co_count': 500},  # 目标检测-计算机视觉
        {'kword1_id': 8, 'kword2_id': 2, 'co_count': 450},  # 图像识别-深度学习
        {'kword1_id': 9, 'kword2_id': 6, 'co_count': 400},  # 人脸识别-计算机视觉
        {'kword1_id': 10, 'kword2_id': 6, 'co_count': 350}, # 视频分析-计算机视觉
        
        # 自然语言处理关联
        {'kword1_id': 11, 'kword2_id': 2, 'co_count': 650}, # 自然语言处理-深度学习
        {'kword1_id': 12, 'kword2_id': 11, 'co_count': 450}, # 机器翻译-自然语言处理
        {'kword1_id': 13, 'kword2_id': 2, 'co_count': 400}, # 语音识别-深度学习
        {'kword1_id': 14, 'kword2_id': 11, 'co_count': 350}, # 文本分类-自然语言处理
        {'kword1_id': 15, 'kword2_id': 11, 'co_count': 300}, # 情感分析-自然语言处理
        
        # 跨领域关联
        {'kword1_id': 1, 'kword2_id': 3, 'co_count': 750},  # 人工智能-机器学习
        {'kword1_id': 6, 'kword2_id': 7, 'co_count': 420},  # 计算机视觉-目标检测
        {'kword1_id': 11, 'kword2_id': 12, 'co_count': 380}, # 自然语言处理-机器翻译
        {'kword1_id': 13, 'kword2_id': 11, 'co_count': 320}, # 语音识别-自然语言处理
        {'kword1_id': 8, 'kword2_id': 9, 'co_count': 280},  # 图像识别-人脸识别
    ]
    
    for data in pairs_data:
        WordPair.objects.create(**data)

def reverse_word_pairs(apps, schema_editor):
    WordPair = apps.get_model('trend', 'WordPair')
    WordPair.objects.all().delete()

class Migration(migrations.Migration):
    dependencies = [
        ('trend', '03_ins_keywords'),
    ]

    operations = [
        migrations.RunPython(insert_word_pairs, reverse_word_pairs),
    ] 